# Direct outputs eligible for public release

# GLOBAL SETTINGS --------------------------------------------------------------

options(
  scipen = 999,
  digits = 16,
  max.print = .Machine$integer.max,
  show.error.locations = TRUE,
  warn = 1
)

RNGkind("L'Ecuyer-CMRG")
seed <- 818675309L
set.seed(seed) # setting main seed

# PACKAGES ---------------------------------------------------------------------
library(data.table)
library(ggplot2)
library(scales)

# PACKAGE SETTINGS -------------------------------------------------------------

# data.table
setDTthreads(threads = 1L)
options(datatable.print.class = TRUE, datatable.print.keys = TRUE)
# so that printing the data.table also shows the variable type on top

# BEGIN FILE -------------------------------------------------------------------

# Read in all estimates

annuitization <-
    readRDS(
        "~/estimation-output/unearned_income_effect_estimates_implied_consumption.rds" #nolint
    )

capitalization <-
    readRDS(
        "~/estimation-output/unearned_income_effect_estimates_implied_consumption_cap_capitalization.rds" #nolint
    )

# Produce Figure 4.1a
# Comparison of estimates of change in n_it across approaches
# (annuitization versus capitalization)
fs <- list()

fs[[1]] <- annuitization[model == "first_stage"]
fs[[1]][, estimate_label := "Annuitization Method"]
fs[[1]][ref_event_time == 0, estimate := NA]
fs[[1]][ref_event_time == 0, cluster_se := NA]

fs[[2]] <- capitalization[model == "first_stage"]
fs[[2]][, estimate_label := "Capitalization Method"]
fs[[2]][ref_event_time == 0, estimate := NA]
fs[[2]][ref_event_time == 0, cluster_se := NA]
fs[[2]][ref_event_time == -1, estimate := NA]
fs[[2]][ref_event_time == -1, cluster_se := NA]

fs <- rbindlist(fs, use.names = TRUE)

# Add omitted
temp <-
    CJ(
        estimate_label = unique(fs$estimate_label),
        model = unique(fs$model),
        income_quartile = unique(fs$income_quartile)
    )
temp[, ref_event_time := -2]
temp[, estimate := 0]
temp[, cluster_se := 0]

fs <- rbindlist(list(fs, temp), use.names = TRUE)
temp <- NULL
rm(temp)

fs[
    ,
    estimate_label :=
        factor(
            estimate_label,
            levels = c("Annuitization Method", "Capitalization Method")
        )
]

figure_4_1_a <-
    ggplot(
        aes(
            x = ref_event_time,
            y = estimate,
            color = factor(estimate_label),
            fill = factor(estimate_label)
        ),
        data = fs
    ) +
    geom_line() +
    geom_ribbon(
        aes(
            ymin = estimate - (1.64 * cluster_se),
            ymax = estimate + (1.64 * cluster_se)
        ),
        linetype = 0,
        alpha = 0.2,
        show.legend = FALSE
    ) +
    theme_bw(base_size = 13) +
    theme(panel.grid.minor = element_blank()) +
    scale_x_continuous(
        breaks = seq.int(from = -7L, to = 5L, by = 1L),
        expand = expansion(mult = c(0.0025, 0.0025))
    ) +
    scale_y_continuous(breaks = seq.int(from = -5000, to = 15000, by = 2500)) +
    labs(x = "Event Time (ℓ)", y = "Event Study Estimate (2016 USD)") +
    theme(
        legend.title = element_blank(),
        legend.position = c(0.17, 0.9175),
        legend.background = element_rect(fill = "white", color = "grey"),
        strip.text.x = element_blank(),
        strip.background = element_blank(),
        legend.key = element_rect(NA),
        legend.spacing.y = unit(2, "mm"),
        legend.key.height = unit(4, "mm"),
        legend.key.width = unit(4, "mm"),
        legend.margin = margin(t = 0, r = 0.1, b = 0.15, l = 0.1, unit = "cm")
    ) +
    scale_color_viridis_d() +
    scale_fill_viridis_d() +
    guides(color = guide_legend(ncol = 1, byrow = TRUE)) +
    coord_cartesian(ylim = c(-5000, 15000))

# need a helper data set
ribbon_dt <- data.table(estimate = seq.int(from = - 50, to = 50, by = 1) / 100)

figure_4_1_a <-
    figure_4_1_a +
    geom_ribbon(
        data = ribbon_dt,
        aes_string(
            x = seq.int(from = - 50, to = 50, by = 1) / 100,
            ymin = -Inf,
            ymax = Inf
        ),
        alpha = 0.1,
        show.legend = FALSE,
        color = NA,
        fill = "black"
    )

ggsave(
    plot = figure_4_1_a,
    filename = "~/paper/figures/figure-4.1a.png",
    width = 6,
    height = 4,
    dpi = 600
)

ggsave(
    plot = figure_4_1_a,
    filename = "~/paper/figures/figure-4.1a.tif",
    width = 6,
    height = 4,
    device = "tiff",
    dpi = 600
)

# Housekeeping
fs <- NULL
ribbon_dt <- NULL
figure_4_1_a <- NULL
rm(fs, ribbon_dt, figure_4_1_a)

# Produce Figure 4.1b
# Comparison of estimates of change in c_it across approaches
# (annuitization versus capitalization)
rf <- list()

rf[[1]] <- annuitization[model == "reduced_form"]
rf[[1]][, estimate_label := "Annuitization Method"]
rf[[1]][ref_event_time == 0, estimate := NA]
rf[[1]][ref_event_time == 0, cluster_se := NA]

rf[[2]] <- capitalization[model == "reduced_form"]
rf[[2]][, estimate_label := "Capitalization Method"]
rf[[2]][ref_event_time == 0, estimate := NA]
rf[[2]][ref_event_time == 0, cluster_se := NA]
rf[[2]][ref_event_time == -1, estimate := NA]
rf[[2]][ref_event_time == -1, cluster_se := NA]

rf <- rbindlist(rf, use.names = TRUE)

# Add omitted
temp <-
    CJ(
        estimate_label = unique(rf$estimate_label),
        model = unique(rf$model),
        income_quartile = unique(rf$income_quartile)
    )
temp[, ref_event_time := -2]
temp[, estimate := 0]
temp[, cluster_se := 0]

rf <- rbindlist(list(rf, temp), use.names = TRUE)
temp <- NULL
rm(temp)

rf[
    ,
    estimate_label :=
        factor(
            estimate_label,
            levels = c("Annuitization Method", "Capitalization Method")
        )
]

figure_4_1_b <-
    ggplot(
        aes(
            x = ref_event_time,
            y = estimate,
            color = factor(estimate_label),
            fill = factor(estimate_label)
        ),
        data = rf
    ) +
    geom_line() +
    geom_ribbon(
        aes(
            ymin = estimate - (1.64 * cluster_se),
            ymax = estimate + (1.64 * cluster_se)
        ),
        linetype = 0,
        alpha = 0.2,
        show.legend = FALSE
    ) +
    theme_bw(base_size = 13) +
    theme(panel.grid.minor = element_blank()) +
    scale_x_continuous(
        breaks = seq.int(from = -7L, to = 5L, by = 1L),
        expand = expansion(mult = c(0.0025, 0.0025))
    ) +
    scale_y_continuous(breaks = seq.int(from = -5000, to = 15000, by = 2500)) +
    labs(x = "Event Time (ℓ)", y = "Event Study Estimate (2016 USD)") +
    theme(
        legend.title = element_blank(),
        legend.position = c(0.17, 0.9175),
        legend.background = element_rect(fill = "white", color = "grey"),
        strip.text.x = element_blank(),
        strip.background = element_blank(),
        legend.key = element_rect(NA),
        legend.spacing.y = unit(2, "mm"),
        legend.key.height = unit(4, "mm"),
        legend.key.width = unit(4, "mm"),
        legend.margin = margin(t = 0, r = 0.1, b = 0.15, l = 0.1, unit = "cm")
    ) +
    scale_color_viridis_d() +
    scale_fill_viridis_d() +
    guides(color = guide_legend(ncol = 1, byrow = TRUE)) +
    coord_cartesian(ylim = c(-5000, 15000))

# need a helper data set
ribbon_dt <- data.table(estimate = seq.int(from = - 50, to = 50, by = 1) / 100)

figure_4_1_b <-
    figure_4_1_b +
    geom_ribbon(
        data = ribbon_dt,
        aes_string(
            x = seq.int(from = - 50, to = 50, by = 1) / 100,
            ymin = -Inf,
            ymax = Inf
        ),
        alpha = 0.1,
        show.legend = FALSE,
        color = NA,
        fill = "black"
    )

ggsave(
    plot = figure_4_1_b,
    filename = "~/paper/figures/figure-4.1b.png",
    width = 6,
    height = 4,
    dpi = 600
)

ggsave(
    plot = figure_4_1_b,
    filename = "~/paper/figures/figure-4.1b.tif",
    width = 6,
    height = 4,
    device = "tiff",
    dpi = 600
)

# Housekeeping
rf <- NULL
ribbon_dt <- NULL
figure_4_1_b <- NULL
annuitization <- NULL
capitalization <- NULL
rm(rf, ribbon_dt, figure_4_1_b, annuitization, capitalization)